'''

将所有股票的历史数据下载到本地, 共有3个参数: 

参数列表
    stock_code 股票代码, 格式为股市类型.股票编码, 如 'sh.603300'
    start_date 开始日期, 格式为 %Y-%m-%d, 如 '2024-01-01', 默认值为'2024-01-01'
    end_date 结束日期, 格式为 %Y-%m-%d, 如 '2024-12-31', 默认值为 None
    forced=False 是否强制写文件，如果保存的文件已存在，且forced为True时，会强制写入。

返回值
    无，文件会被保存在 f'./daily_data/{stock_code}.csv', 如 ./daily_data/sh.600749.csv。
    
使用示例
    download_stock(stock_code, '2000-01-01')  # 下载华铁应急

【参考资料】
1. 技术文档地址: https://ai-cyber-security.feishu.cn/docx/AB1fdqV9xo9K9axwT3hcMOOcnzc
2. BaoStock官网: http://baostock.com/baostock/index.php/Python_API文档

'''
import os, sys
import baostock as bs
from datetime import datetime, timedelta
import pandas as pd
from utils.common import *

print('功能说明: 下载指定股票或所有股票。\n')
print('函数定义:\n    python ./download_daily_data.py [stock_code]\n')
print('    stock_code     更新指定股票代码，要以sh.或sz.开头, 如sz.000430')
print('使用示例:\n    python ./download_daily_data.py sz.000430\n')



# 读取数据, 默认开始日期为2000/01/01, 结束日期为当前日期
def download_stock(lg, stock_code, start_date='2000-01-01', end_date=None, out_dir='daily_data', forced=False):
    stock_file = os.path.join(out_dir, f'{stock_code}.csv')
    if os.path.exists(stock_file) and not forced:
        return None
    
    # 初始化参数
    end_date= datetime.now().strftime('%Y-%m-%d') if end_date == None else end_date
    
    # 读取数据 column_names = "date,preclose,open,high,low,close"
    column_names = "date,code,open,high,low,close,preclose,volume,amount,adjustflag,turn,tradestatus,pctChg,isST"
    rs = bs.query_history_k_data_plus(stock_code, column_names, 
         start_date=start_date, end_date=end_date, frequency="d", adjustflag="3")

    data_list = []
    if rs.error_code == '0':
        while (rs.error_code == '0') & rs.next():
            data_list.append(rs.get_row_data())
    else:
        print('query_history_k_data_plus failed, error_code: ' + rs.error_code)
        print('query_history_k_data_plus failed, error_msg: ' + rs.error_msg)
        return None
    
    # 保存到本地, 并返回df
    df = pd.DataFrame(data_list, columns=column_names.split(','))
    df.to_csv(stock_file, encoding="utf-8", index=False)
    return df


def download_all_stocks(date='2024-12-02'):
    ''' 注意: 节假日返回为0 '''
    lg = bs.login()
    if  lg.error_code != '0':
        print("Loging failed. Error code: " + lg.error_msg)
        print("Loging failed. Error message: " + lg.error_code)
        return
    
    # 获取指定日期的指数、股票数据
    stock_rs = bs.query_all_stock(date)
    stock_df = stock_rs.get_data()
    stock_df.to_csv('allstocks.csv')
    print('all stocks saved to: allstocks.csv, size:', len(stock_df))
    return

# 如果不存在, 则下载所有股票代码
single_stock = ''
if len(sys.argv) > 1:
    single_stock = sys.argv[1]

stocknames_path = 'all_stock_codes.csv'
if not os.path.exists(stocknames_path):
    print(f'stock name file "{stocknames_path}" does not exist.')
    sys.exit()

data_dir = 'daily_data'
if not os.path.exists(stocknames_path):
    download_all_stocks()

if not os.path.exists(data_dir):
    os.mkdir(data_dir)

lg = bs.login()
if  lg.error_code != '0':
    print("Loging failed. Error code: " + lg.error_msg)
    print("Loging failed. Error message: " + lg.error_code)
    sys.exit()

# 用于控制输出显示的宽度
code = 'SH.603538' if len(sys.argv) < 2 else sys.argv[1].upper()
start_date = '2010-01-01' if len(sys.argv) < 3 else sys.argv[2]

# 文件包括三列:  code, tradeStatus, code_name, 如：sz.301589 1 诺瓦星云
df_codes = pd.read_csv(stocknames_path)
start_date = '2025-01-01'
end_date = f'{datetime.now():%Y-%m-%d}'
for index, row in df_codes.iterrows():
    down_res = 'downloaded'
    dr = row.tolist()
    t1 = datetime.now()
    if len(single_stock) > 0 and single_stock == dr[0]:
        download_stock(lg, dr[0], start_date, end_date, data_dir, forced=True)
    else:
        res = download_stock(lg, dr[0], start_date, end_date, data_dir)
        if res is None:
            down_res = 'existed'
    t2 = datetime.now()
    print(f'{index:4}/{len(df_codes)} {(t2-t1).total_seconds():.2f} {dr[0]}  {dr[2]} \t {down_res}.', flush=True)
#    sys.stdout.flush()
bs.logout()
